from shutil import copyfile
from video_compressor import THREAD_POOL
from video_compressor import printf
from video_compressor import Colors
from video_compressor import VIDEO_WIDE
from video_compressor import VIDEO_HIGH
from video_compressor import IN_DIR_PATH
from video_compressor import OUT_DIR_PATH
from video_compressor import VIDEO_EXTS
from video_compressor import MIX_TIME
from video_compressor import MAX_FRAMES
from video_compressor import get_file
from video_compressor import split_duration_by_count
from video_compressor.utils import get_video_info
from video_compressor.ffmpeg_video import ffmpy_gpu_video_slice
from video_compressor.ffmpeg_video import run


def get_tasks():
    tasks, file_set = [], set()
    file_iter = get_file(IN_DIR_PATH)
    for file in file_iter:
        out_dir_path = OUT_DIR_PATH.joinpath(file.relative_to(IN_DIR_PATH))
        if not out_dir_path.parent.is_dir(): out_dir_path.parent.mkdir(parents=True, exist_ok=True)
        if out_dir_path.suffix not in VIDEO_EXTS:
            copyfile(file, out_dir_path)
        else:
            stem_list = file.stem.split('-')
            if len(stem_list) == 1 and stem_list[0].isdigit(): copyfile(file, out_dir_path)
            else: file_set.add(file)

    for file in file_set:
        out_dir_path = OUT_DIR_PATH.joinpath(file.relative_to(IN_DIR_PATH))
        stem_list = file.stem.split('-')
        wide, high, total_time = get_video_info(file)
        video_wide = VIDEO_WIDE if VIDEO_WIDE else wide
        video_high = VIDEO_HIGH if VIDEO_HIGH else high
        if len(stem_list) == 2 and stem_list[0].isdigit() and stem_list[1].isdigit():
            stem = list(map(lambda x: int(x), stem_list))
            min_num, max_num = min(stem), max(stem)
            slice_lier = split_duration_by_count(total_time, (max_num - min_num + 1))
        elif not stem_list[0].isdigit():
            min_num = 1
            num_segments = min([int(total_time // MIX_TIME), MAX_FRAMES])
            slice_lier = split_duration_by_count(total_time, num_segments)
        else:
            min_num, slice_lier = 0, []
        for i, (start, end) in enumerate(slice_lier, start=min_num):
            out_path = out_dir_path.parent / f'{i:03d}{out_dir_path.suffix}'
            if not out_path.is_file():
                printf(f'{Colors.OKGREEN}创建任务{Colors.ENDC}{Colors.OKBLUE}{out_path}{Colors.ENDC} {Colors.OKGREEN}分割时长：{end - start}{Colors.ENDC}')
                ffmpeg = ffmpy_gpu_video_slice(file, out_path, start, end, video_wide, video_high)
                tasks.append(THREAD_POOL.submit(run, ffmpeg, out_path))
    return tasks
